package com.dahuyou.method.record;

import com.dahuyou.method.record.filter.MethodRecordingFilter;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

public class MethodRecordingTransformer implements ClassFileTransformer {

    /**
     * 所有的类被来加载器加载的时候都会走这个方法，这样我们就有机会通过各种技术手段来对字节码进行插桩了
     *
     * @param loader
     * @param className           类全限定名称
     * @param classBeingRedefined
     * @param protectionDomain
     * @param classfileBuffer     类字节码对应的二进制数组，可以通过defineClass直接加载到JVM中并生成Class对象
     * @return
     * @throws IllegalClassFormatException
     */
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
            // 过滤不需要插桩的类
            if (MethodRecordingFilter.isNotNeedInject(className)) {
                return classfileBuffer;
            }
            // 插桩
            byte[] bytes = generateInstrumentationCode(loader, className, classfileBuffer);
            // 写到磁盘的.class文件中，方便调试
            outputClazz(bytes, className);
            return bytes;
        } catch (Throwable e) {
            System.out.println(e.getMessage());
        }
        return classfileBuffer;

    }

    private void outputClazz(byte[] bytes, String className) {
        // 输出类字节码
        FileOutputStream out = null;
        try {
            String pathName = MethodRecordingTransformer.class.getResource("/").getPath() + className + "SQM.class";
            out = new FileOutputStream(new File(pathName));
            System.out.println("插桩后代码输出路径：" + pathName);
            out.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (null != out) try {
                out.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 生成插桩代码
     *
     * @param loader
     * @param className
     * @param classfileBuffer
     * @return
     */
    private byte[] generateInstrumentationCode(ClassLoader loader, String className, byte[] classfileBuffer) {
        // 读取原有类
        ClassReader cr = new ClassReader(classfileBuffer);
        // 通过writer向需要插桩类写入插桩逻辑
        ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_MAXS);
        // 使用子自定义的类访问器来定义插桩逻辑
        ClassVisitor cv = new MethodRecordingClassAdapter(cw, className);
        // 真正插桩
        cr.accept(cv, ClassReader.EXPAND_FRAMES);
        return cw.toByteArray();
    }
}
